{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "589389e6-2aad-45b6-bc03-e9b37ab1514f",
   "metadata": {},
   "source": [
    "# Reinforcement learning for legged robots\n",
    "\n",
    "## Setup\n",
    "\n",
    "Before we start, you will need to update your conda environment to use Gymnasium (maintained) rather than OpenAI Gym (discontinued). You can simply run:\n",
    "\n",
    "```\n",
    "conda env create -f robotics_rl_env.yml\n",
    "conda activate robotics_course_rl\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3fa7b15-843f-4696-9bfc-141da71bf7d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "import stable_baselines3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e264559-88cc-48b7-8257-f7755fff3ce7",
   "metadata": {},
   "source": [
    "Let's import the usual suspects as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86303cf2-f879-407d-b528-6c0a80b8df20",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pylab as plt\n",
    "import numpy as np\n",
    "\n",
    "plt.ion()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b09e16a1-bf4a-403b-8709-3da11bc3c4b4",
   "metadata": {},
   "source": [
    "# Inverted pendulum environment\n",
    "\n",
    "The inverted pendulum model is not just a toy model reproducing the properties of real robot models for balancing: as it turns out, the inverted pendulum appears in the dynamics of *any* mobile robot, that is, a model with a floating-base joint at the root of the kinematic tree. (If you are curious: the inverted pendulum is a limit case of the [Newton-Euler equations](https://scaron.info/robotics/newton-euler-equations.html) corresponding to floating-base coordinates in the equations of motion $M \\ddot{q} + h = S^T \\tau + J_c^T f$, in the limit where the robot [does not vary its angular momentum](https://scaron.info/robotics/point-mass-model.html).) Thus, while we work on a simplified inverted pendulum in this notebook, concepts and tools are those used as-is on real robots, as you can verify by exploring the bonus section.\n",
    "\n",
    "Gymnasium is mainly a single-agent reinforcement learning API, but it also comes with simple environments, including an inverted pendulum sliding on a linear guide:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11b5d942-85fa-435e-b5ef-8c85a74ba3db",
   "metadata": {},
   "outputs": [],
   "source": [
    "with gym.make(\"InvertedPendulum-v4\", render_mode=\"human\") as env:\n",
    "    action = 0.0 * env.action_space.sample()\n",
    "    observation, _ = env.reset()\n",
    "    episode_return = 0.0\n",
    "    for step in range(200):\n",
    "        # action[0] = 5.0 * observation[1] + 0.3 * observation[0]\n",
    "        observation, reward, terminated, truncated, _ = env.step(action)\n",
    "        episode_return += reward\n",
    "        if terminated or truncated:\n",
    "            observation, _ = env.reset()\n",
    "    print(f\"Return of the episode: {episode_return}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7322422-94db-4e12-b299-36bb40649cf7",
   "metadata": {},
   "source": [
    "The structure of the action and observation vectors are documented in [Inverted Pendulum - Gymnasium Documentation](https://gymnasium.farama.org/environments/mujoco/inverted_pendulum/).  The observation, in particular, is a NumPy array with four coordinates that we recall here for reference:\n",
    "\n",
    "| Num | Observation | Min | Max | Unit |\n",
    "|-----|-------------|-----|-----|------|\n",
    "|   0 | position of the cart along the linear surface | -Inf | Inf | position (m) |\n",
    "|   1 | vertical angle of the pole on the cart | -Inf | Inf | angle (rad) |\n",
    "|   2 | linear velocity of the cart | -Inf | Inf | linear velocity (m/s) |\n",
    "|   3 | angular velocity of the pole on the cart | -Inf | Inf | anglular velocity (rad/s) |\n",
    "\n",
    "We will use the following labels to annotate plots:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3231c70-f49d-49be-b260-aadbade7b403",
   "metadata": {},
   "outputs": [],
   "source": [
    "OBSERVATION_LEGEND = (\"pitch\", \"position\", \"linear_velocity\", \"angular_velocity\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa062536-204c-4312-a858-f992f3db61d6",
   "metadata": {},
   "source": [
    "Check out the documentation for the definitions of the action and rewards."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c285d7ce-3a97-4b07-8b5f-a9b04d7721ab",
   "metadata": {},
   "source": [
    "# PID control\n",
    "\n",
    "A *massively* used class of policies is the [PID controller](https://en.wikipedia.org/wiki/Proportional%E2%80%93integral%E2%80%93derivative_controller). Let's say we have a reference observation, like $o^* = [0\\ 0\\ 0\\ 0]$ for the inverted pendulum. Denoting by $e(t) = o^* - o(t)$ the *error* of the system when it observes a given state, a continuous-time PID controller will apply the action:\n",
    "\n",
    "$$\n",
    "a(t) = K_p^T e(t) + K_d^T \\dot{e}(t) + K_i^T \\int e(\\tau) \\mathrm{d} \\tau\n",
    "$$\n",
    "\n",
    "where $K_{p}, K_i, K_d \\in \\mathbb{R}^4$ are constants called *gains* and tuned by the user. In discrete time the idea is the same:\n",
    "\n",
    "$$\n",
    "a_k = K_p^T e_k + K_d^T \\frac{e_k - e_{k-1}}{\\delta t} + K_i^T \\sum_{i=0}^{k} e_i {\\delta t}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63c381eb-fca9-4ef4-8f99-3b1943231654",
   "metadata": {},
   "source": [
    "Let's refactor the rolling out of our episode into a standalone function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c839bc6-168a-42c3-8f1c-c6b0c5411901",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rollout_from_env(env, policy):\n",
    "    episode = []\n",
    "    observation, _ = env.reset()\n",
    "    episode.append(observation)\n",
    "    for step in range(1000):\n",
    "        action = policy(observation)\n",
    "        observation, reward, terminated, truncated, _ = env.step(action)\n",
    "        episode.extend([action, reward, observation])\n",
    "        if terminated or truncated:\n",
    "            return episode\n",
    "    return episode\n",
    "\n",
    "def rollout(policy, show: bool = True):\n",
    "    kwargs = {\"render_mode\": \"human\"} if show else {}\n",
    "    with gym.make(\"InvertedPendulum-v4\", **kwargs) as env:\n",
    "        episode = rollout_from_env(env, policy)\n",
    "    return episode"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79ff0dce-a4df-4917-bb17-2393353610a3",
   "metadata": {},
   "source": [
    "## Question 1: Write a PID controller that balances the inverted pendulum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7cfb28b-ff73-42ff-9524-eac8ec12f8a1",
   "metadata": {},
   "source": [
    "You can use global variables to store the (discrete) derivative and integral terms, this will be OK here as we only rollout a single trajectory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "045ddcef-c0f7-4251-b73f-d5df5a0027e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pid_policy(observation: np.ndarray) -> np.ndarray:\n",
    "    my_action_value: float = 0.0  # your action here\n",
    "    return np.array([my_action_value])\n",
    "\n",
    "episode = rollout(pid_policy, show=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0a005aa-87fa-4f98-8ace-f24421886bed",
   "metadata": {},
   "source": [
    "You can look at the system using `show=True`, but intuition usually builds faster when looking at relevant plots:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9aa5decd-779c-4f0d-84fd-3eb47358b7fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "observations = np.array(episode[::3])\n",
    "\n",
    "plt.plot(observations)\n",
    "plt.legend(OBSERVATION_LEGEND)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98d50cd2-26fa-4d3c-a671-1ed0e1b9ee93",
   "metadata": {},
   "source": [
    "Can you reach the full reward of 1000 steps?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bacbd0a-2ac5-44cf-848b-8ebfb6fe35d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Return of the episode: {sum(episode[2::3])}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b17cc998-1b23-416f-8e3b-810100c223fb",
   "metadata": {},
   "source": [
    "# Policy optimization\n",
    "\n",
    "Let us now train a policy, parameterized by a multilayer perceptron (MLP), to maximize the expected return over episodes on the inverted pendulum environment."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5631f0f-1b84-4ee6-8e9c-b4f2915bd281",
   "metadata": {},
   "source": [
    "## Our very first policy\n",
    "\n",
    "We will use the proximal policy optimization (PPO) algorithm for training, using the implementation from Stable Baselines3: [PPO - Stable Baselines3 documentation](https://stable-baselines3.readthedocs.io/en/master/modules/ppo.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "128867ca-e600-4ba1-abbd-1f918976fba2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from stable_baselines3 import PPO\n",
    "\n",
    "with gym.make(\"InvertedPendulum-v4\", render_mode=\"human\") as env:\n",
    "    first_policy = PPO(\"MlpPolicy\", env, verbose=0)\n",
    "    first_policy.learn(total_timesteps=1000, progress_bar=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6323400b-18ca-43f6-a81e-a5e7f033a536",
   "metadata": {},
   "source": [
    "By instantiating the algorithm with no further ado, we let the library decide for us on a sane set of default hyperparameters, including:\n",
    "\n",
    "- Rollout buffers of `n_steps = 2048` steps, which we will visit `n_epochs = 10` times with mini-batches of size `batch_size = 64`.\n",
    "- Clipping range: ``0.2``.\n",
    "- No entropy regularization.\n",
    "- Learning rate for the Adam optimizer: ``3e-4``\n",
    "- Policy and value-function network architectures: two layers of 64 neurons with $\\tanh$ activation functions.\n",
    "\n",
    "We then called the `learn` function to execute PPO over a fixed total number of timesteps, here just a thousand."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b82173c-6609-4b83-8618-36f82c1c1373",
   "metadata": {},
   "source": [
    "Rendering actually took a significant chunk of time. Let's instantiate and keep an environment open without rendering:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "460fe1c7-ee3b-450a-b09c-03b96f9086bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(\"InvertedPendulum-v4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9bd090f-ca34-41e0-9900-52977eef9c4b",
   "metadata": {},
   "source": [
    "We can use it to train much more steps in roughly the same time, reporting training metrics every `n_steps` step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7262602-c277-4697-8987-ba126a87e75b",
   "metadata": {},
   "outputs": [],
   "source": [
    "second_policy = PPO(\"MlpPolicy\", env, verbose=1)\n",
    "second_policy.learn(total_timesteps=10_000, progress_bar=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6219aab8-1143-4606-a44f-b62fdffebbf1",
   "metadata": {},
   "source": [
    "Let's see how this policy performs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f17dc178-bb9c-4155-8047-feed1e575226",
   "metadata": {},
   "outputs": [],
   "source": [
    "def policy_closure(policy):\n",
    "    \"\"\"Utility function to turn our policy instance into a function.\n",
    "\n",
    "    Args:\n",
    "        policy: Policy to turn into a function.\n",
    "        \n",
    "    Returns:\n",
    "        Function from observation to policy action.\n",
    "    \"\"\"\n",
    "    def policy_function(observation):\n",
    "        action, _ = policy.predict(observation)\n",
    "        return action\n",
    "    return policy_function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e4e3cd4-4572-4c40-94b1-688d472a4b8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "episode = rollout(policy_closure(second_policy), show=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "941219f1-9b3c-4e66-86e5-d42f2473b149",
   "metadata": {},
   "source": [
    "Okay, it looks like we didn't train for long enough!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "195f2a85-8dd7-4f3a-8368-7427f1caadca",
   "metadata": {},
   "source": [
    "## Monitoring performance during training\n",
    "\n",
    "Let's train for longer, and use TensorBoard to keep track. We don't know how long training will take so let's put a rather large total number of steps (you can interrupt training once you observed convergence in TensorBoard):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "196da0ad-1e83-441e-ac10-6c9ecd83c224",
   "metadata": {},
   "outputs": [],
   "source": [
    "erudite_policy = PPO(\n",
    "    \"MlpPolicy\",\n",
    "    env,\n",
    "    tensorboard_log=\"./inverted_pendulum_tensorboard/\",\n",
    "    verbose=0,\n",
    ")\n",
    "\n",
    "erudite_policy.learn(\n",
    "    total_timesteps=1_000_000,\n",
    "    progress_bar=False,\n",
    "    tb_log_name=\"erudite\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad91d14e-53e2-443f-b7ab-edd69d480add",
   "metadata": {},
   "source": [
    "Run TensorBoard on the directory thus created to open a dashboard in your Web browser:\n",
    "\n",
    "```\n",
    "tensorboard --logdir ./inverted_pendulum_tensorboard/\n",
    "```\n",
    "\n",
    "The link will typically be http://localhost:6006 (port number increases if you run TensorBoard multiple times in parallel). Tips:\n",
    "\n",
    "- Click the Settings icon in the top-right corner and enable \"Reload data\"\n",
    "- Uncheck \"Ignore outliers in chart scaling\" (your preference)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68771e35-48cd-43be-89ff-0055dc196d0b",
   "metadata": {},
   "source": [
    "## Saving our policy\n",
    "\n",
    "Now that you spent some computing to optimize an actual policy, better save it to disk:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "effeccfc-8b95-48e1-98c4-1b96838bb28e",
   "metadata": {},
   "outputs": [],
   "source": [
    "erudite_policy.save(\"pendulum_erudite\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea09c56e-8647-414d-aae0-5e1b16ba3a0f",
   "metadata": {},
   "source": [
    "You can then reload it later by:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "994dae4f-b651-4488-925e-2ba369eeedc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "erudite_policy = PPO.load(\"pendulum_erudite\", env=env)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dded4f1e-ae57-4c94-b2e1-3a6d019aecc4",
   "metadata": {},
   "source": [
    "## Question 2: How many steps does it take to train a successful policy?\n",
    "\n",
    "We consider a policy successful if it consistently achieves the maximum return of 1000."
   ]
  },
  {
   "cell_type": "raw",
   "id": "42c6d68d-4812-4222-97da-a6699803b986",
   "metadata": {},
   "source": [
    "== Your reply here =="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "553b846f-db13-43ba-81cb-57b039852c86",
   "metadata": {},
   "source": [
    "## A more realistic environment\n",
    "\n",
    "Real systems suffer from the two main issues we saw in the [Perception and estimation](https://scaron.info/robotics-mva/#5-perception-estimation) class: *bias* and *variance*. In this section, we model bias in actuation and perception by adding delays (via low-pass filtering) to respectively the action and observation vectors. Empirically this is an effective model, as for instance it contributes to sim2real transfer on Upkie. To add these delays, we use an [`environment wrapper`](https://gymnasium.farama.org/api/wrappers/), which is a convenient way to compose environments, used in both the Gymnasium and Stable Baselines3 APIs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e8a3140-7ee7-4d6f-afd9-19d6ca4816c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DelayWrapper(gym.Wrapper):\n",
    "    def __init__(self, env, time_constant: float = 0.2):\n",
    "        \"\"\"Wrap environment with some actuation and perception modeling.\n",
    "\n",
    "        Args:\n",
    "            env: Environment to wrap.\n",
    "            time_constant: Constant of the internal low-pass filter, in seconds.\n",
    "                Feel free to play with different values but leave it to the default\n",
    "                of 0.2 seconds when handing out your homework.\n",
    "\n",
    "        Note:\n",
    "            Delays are implemented by a low-pass filter. The same time constant\n",
    "            is used for both actions and observations, which is not realistic, but\n",
    "            makes for less tutorial code ;)\n",
    "        \"\"\"\n",
    "        alpha = env.dt / time_constant\n",
    "        assert 0.0 < alpha < 1.0\n",
    "        super().__init__(env)\n",
    "        self._alpha = alpha\n",
    "        self._prev_action = 0.0 * env.action_space.sample()\n",
    "        self._prev_observation = np.zeros(4)\n",
    "\n",
    "    def low_pass_filter(self, old_value, new_value):\n",
    "        return old_value + self._alpha * (new_value - old_value)\n",
    "        \n",
    "    def step(self, action):\n",
    "        new_action = self.low_pass_filter(self._prev_action, action)\n",
    "        observation, reward, terminated, truncated, info = self.env.step(new_action)\n",
    "        new_observation = self.low_pass_filter(self._prev_observation, observation)\n",
    "        self._prev_action = new_action\n",
    "        self._prev_observation = new_observation\n",
    "        return new_observation, reward, terminated, truncated, info\n",
    "\n",
    "delay_env = DelayWrapper(env)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1b5de5e-50ca-4049-bb5f-b9203919e0ba",
   "metadata": {},
   "source": [
    "We can check how our current policy fares against the delayed environment. Spoiler alert: no great."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e1508e6-e04f-4b22-8009-80baae1bae7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "delay_episode = rollout_from_env(delay_env, policy_closure(erudite_policy))\n",
    "observations = np.array(delay_episode[::3])\n",
    "\n",
    "plt.plot(observations[:, :2])\n",
    "plt.legend(OBSERVATION_LEGEND)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70af3932-751e-47e4-8334-bd55be62aaa1",
   "metadata": {},
   "source": [
    "## Question 3: Can't we just re-train a policy on the new environment?\n",
    "\n",
    "At this point of the tutorial this is a rethorical question, but we should check anyway. Re-train a policy on the delayed environment and comment on its performance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "693aa97c-3ee2-4cbd-bc06-7cb224e8bc86",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your code here"
   ]
  },
  {
   "cell_type": "raw",
   "id": "48928906-bcd9-40d5-b17e-35fd06d6c6ac",
   "metadata": {},
   "source": [
    "== Your observations here =="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0e2df30-259f-477a-ab14-d39c17e5f15f",
   "metadata": {},
   "source": [
    "## The Real Question 3: Why do delays degrade both runtime and training performance?\n",
    "\n",
    "Loss in runtime performance refers to the one we observed when executing a policy trained without delay on an environment with delays. Loss in training performance refers to the fact that, even when we train a new policy on the environment with delays, by the end of training it does not achieve maximum return."
   ]
  },
  {
   "cell_type": "raw",
   "id": "3b7459d5-93d0-49cb-85c7-2172e2b08073",
   "metadata": {},
   "source": [
    "== Your explanation here =="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e63a441a-a84d-49ab-aecc-7362dee66b91",
   "metadata": {},
   "source": [
    "Propose and implement a way to overcome this. Train the resulting policy in a variable called `iron_policy`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b22770ba-4e58-4989-b62c-d5aa1734336c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a7a876f-e78e-47bb-9b42-9423618d1e42",
   "metadata": {},
   "outputs": [],
   "source": [
    "iron_policy.save(\"iron_policy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2a70b63-7fda-4c0f-b777-ef0dc2128ab2",
   "metadata": {},
   "source": [
    "Roll out an episode and plot the outcome to show that your policy handles delays properly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45f13a2b-6a1c-4d44-bb84-2fffaf6bf6e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your episode rollout here\n",
    "\n",
    "plt.plot(np.array(observations)[:, :2])\n",
    "plt.legend(OBSERVATION_LEGEND)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e12fcf1-88b9-4899-b79d-866c67e4a3f5",
   "metadata": {},
   "source": [
    "## Question 4: Can you improve sampling efficiency?\n",
    "\n",
    "This last question is open: what can you change in the pipeline to train a policy that achieves maximum return using less samples? Report on at least one thing that allowed you to train with less environment steps."
   ]
  },
  {
   "cell_type": "raw",
   "id": "0f5cb9a5-fd18-4077-a6fc-83fa5377de96",
   "metadata": {},
   "source": [
    "== Your report here =="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "131966f5-9524-4b44-9843-0c1a662ba2e1",
   "metadata": {},
   "source": [
    "Here is a state-of-the-art™ utility function if you want to experiment with scheduling some of the ``Callable[[float], float]`` [hyperparameters](https://stable-baselines3.readthedocs.io/en/master/modules/ppo.html#parameters):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3de11ab9-2534-4723-8868-1582772d038c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def affine_schedule(y_0: float, y_1: float):\n",
    "    \"\"\"Affine schedule as a function over the [0, 1] interval.\n",
    "\n",
    "    Args:\n",
    "        y_0: Function value at zero.\n",
    "        y_1: Function value at one.\n",
    "        \n",
    "    Returns:\n",
    "        Corresponding affine function.\n",
    "    \"\"\"\n",
    "    def schedule(x: float) -> float:\n",
    "        return y_0 + x * (y_1 - y_0)\n",
    "    return schedule"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b21d78dd-f80e-4183-8fa7-55c803e38404",
   "metadata": {},
   "source": [
    "And here is a wrapper template if you want to experiment with reward shaping:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cf9a3ed-8f76-4fac-98f2-3a23df818deb",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomRewardWrapper(gym.Wrapper):\n",
    "   def __init__(self, env):\n",
    "       super().__init__(env)\n",
    "\n",
    "   def step(self, action):\n",
    "       observation, reward, terminated, truncated, info = self.env.step(action)\n",
    "       new_reward = 0.0  # your formula here\n",
    "       return observation, new_reward, terminated, truncated, info"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c4dd0df-6dc7-4d51-b29b-c77b49bde437",
   "metadata": {},
   "source": [
    "# Bonus: training a policy for a real robot\n",
    "\n",
    "This section is entirely optional and will only work on Linux or macOS. In this part, we follow the same training pipeline but with the open source software of [Upkie](https://hackaday.io/project/185729-upkie-wheeled-biped-robots)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9634ff93-f09f-4e0a-8d0f-547848f3900b",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "<img src=\"https://user-images.githubusercontent.com/1189580/170496331-e1293dd3-b50c-40ee-9c2e-f75f3096ebd8.png\" style=\"height: 100px\" align=\"right\" />\n",
    "\n",
    "First, make sure you have a C++ compiler (setup one-liners: [Fedora](https://github.com/upkie/upkie/discussions/100), [Ubuntu](https://github.com/upkie/upkie/discussions/101)). You can run an Upkie simulation right from the command line. It won't install anything on your machine, everything will run locally from the repository:\n",
    "\n",
    "```console\n",
    "git clone https://github.com/upkie/upkie.git\n",
    "cd upkie\n",
    "git checkout fb9a0ab1f67a8014c08b34d7c0d317c7a8f71662\n",
    "./start_simulation.sh\n",
    "cd ..\n",
    "pip install upkie\n",
    "```\n",
    "\n",
    "**NB:** this tutorial is written for the specific commit checked out above. If some instructions don't work it's likely you forgot to check it out."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba44abc0-f7e9-4c2b-9d4e-a3579213e138",
   "metadata": {},
   "source": [
    "## Stepping the environment\n",
    "\n",
    "If everything worked well, you should be able to step an environment as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acedf0d6-fc2f-43f4-9ff6-a8e12dbd7ae0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "import upkie.envs\n",
    "\n",
    "upkie.envs.register()\n",
    "\n",
    "episode_return = 0.0\n",
    "with gym.make(\"UpkieGroundVelocity-v1\", frequency=200.0) as env:\n",
    "    observation, _ = env.reset()  # connects to the spine (simulator or real robot)\n",
    "    action = 0.0 * env.action_space.sample()\n",
    "    for step in range(1000):\n",
    "        pitch = observation[0]\n",
    "        action[0] = 10.0 * pitch  # 1D action: [ground_velocity]\n",
    "        observation, reward, terminated, truncated, _ = env.step(action)\n",
    "        episode_return += reward\n",
    "        if terminated or truncated:\n",
    "            observation, _ = env.reset()\n",
    "\n",
    "print(f\"We have stepped the environment {step + 1} times\")\n",
    "print(f\"The return of our episode is {episode_return}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "031343b5-cf94-46ae-98f3-a4c5ebbc037c",
   "metadata": {},
   "source": [
    "(If you see a message \"Waiting for spine /vulp to start\", it means the simulation is not running.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aecfd91f-676c-4d6f-beb0-a286dc681ae3",
   "metadata": {},
   "source": [
    "We can double-check the last observation from the episode:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed6d972f-4cc9-4005-b9a1-4a7433a19938",
   "metadata": {},
   "outputs": [],
   "source": [
    "def report_last_observation(observation):\n",
    "    print(\"The last observation of the episode is:\")\n",
    "    print(f\"- Pitch from torso to world: {observation[0]:.2} rad\")\n",
    "    print(f\"- Ground position: {observation[1]:.2} m\")\n",
    "    print(f\"- Angular velocity from torso to world in torso: {observation[2]:.2} rad/s\")\n",
    "    print(f\"- Ground velocity: {observation[3]:.2} m/s\")\n",
    "    \n",
    "report_last_observation(observation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5a269e3-d876-4d05-88e5-b0d73be6f939",
   "metadata": {},
   "source": [
    "## Question B1: PID control\n",
    "\n",
    "Adapt your code from Question 1 to this environment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e8e9f1e-f2a1-4a18-aa38-256a425d018c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def policy_b1(observation):\n",
    "    return np.array([0.0])  # replace with your solution\n",
    "\n",
    "\n",
    "def run(policy, nb_steps: int):\n",
    "    episode_return = 0.0\n",
    "    with gym.make(\"UpkieGroundVelocity-v1\", frequency=200.0) as env:\n",
    "        observation, _ = env.reset()  # connects to the spine (simulator or real robot)\n",
    "        for step in range(nb_steps):\n",
    "            action = policy_b1(observation)\n",
    "            observation, reward, terminated, truncated, _ = env.step(action)\n",
    "            if terminated or truncated:\n",
    "                print(\"Fall detected!\")\n",
    "                return episode_return\n",
    "    report_last_observation(observation)\n",
    "    return episode_return\n",
    "\n",
    "\n",
    "episode_return = run(policy_b1, 1000)\n",
    "print(f\"The return of our episode is {episode_return}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e999eb22-a94a-4a58-ac06-9a7dbc15a7ee",
   "metadata": {},
   "source": [
    "## Training a new policy\n",
    "\n",
    "The Upkie repository ships three agents based on PID control, model predictive control and reinforcement learning. We now focus on the latter, called the \"PPO balancer\".\n",
    "\n",
    "Check that you can run the training part by running, from the root of the repository:\n",
    "\n",
    "```\n",
    "./tools/bazel run //agents/ppo_balancer:train -- --nb-envs 1 --show\n",
    "```\n",
    "\n",
    "A simulation window should pop, and verbose output from SB3 should be printed to your terminal.\n",
    "\n",
    "By default, training data will be logged to `/tmp`. You can select a different output path by setting the `UPKIE_TRAINING_PATH` environment variable in your shell. For instance:\n",
    "\n",
    "```\n",
    "export UPKIE_TRAINING_PATH=\"${HOME}/src/upkie/training\"\n",
    "```\n",
    "\n",
    "Run TensorBoard from the training directory:\n",
    "\n",
    "```\n",
    "tensorboard --logdir ${UPKIE_TRAINING_PATH}  # or /tmp if you keep the default\n",
    "```\n",
    "\n",
    "Each training will be named after a word picked at random in an English dictionary."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7e47aad-7787-409c-af7e-b83bfccaa592",
   "metadata": {},
   "source": [
    "## Selecting the number of processes\n",
    "\n",
    "We can increase the number of parallel CPU environments ``--nb-envs`` to a value suitable to your computer. Let training run for a minute and check `time/fps`. Increase the number of environments and compare the stationary regime of `time/fps`. You should see a performance increase when adding the first few environments, followed by a declined when there are two many parallel processes compared to your number of CPU cores. Pick the value that works best for you."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "696a0943-cc10-4cd0-a2d8-d5313dbe37e5",
   "metadata": {},
   "source": [
    "## Running a trained policy\n",
    "\n",
    "Copy the file `final.zip` from your trained policy directory to `agents/ppo_balancer/policy/params.zip`. Start a simulation and run the policy by:\n",
    "\n",
    "```\n",
    "./tools/bazel run //agents/ppo_balancer\n",
    "```\n",
    "\n",
    "What kind of behavior do you observe?"
   ]
  },
  {
   "cell_type": "raw",
   "id": "eaabe73c-f412-44b5-a714-241077720d01",
   "metadata": {},
   "source": [
    "== Your observations here =="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c356c81-b5ef-4364-a5db-c8e2600e104a",
   "metadata": {},
   "source": [
    "## Question B2: Improve this baseline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "527ecb8c-7292-432f-b0d3-b90c36de8719",
   "metadata": {},
   "source": [
    "The policy you are testing here is not the one we saw in class. Open question: improve on it using any of the methods we discussed. Measure the improvement by `ep_len_mean` or any other quantitative criterion:"
   ]
  },
  {
   "cell_type": "raw",
   "id": "ce7d720b-17b8-493d-8128-e66c6571d3ff",
   "metadata": {},
   "source": [
    "== Your experiments here ==\n",
    "\n",
    "- Tried: ... / Measured outcome: ..."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
